###################################################################
# HIGH LEVEL FUNCTIONS TO DISTRIBUTE R ACROSS NODES IN A CLUSTER
# RMPI (R WRAPPER LIBRARY FOR MPI PRIMITIVES) IS USED

# load relevant source files
source("timing.R")
source("commonFuncs.R")

####################################################################
# Cluster creation / deletion
#
startCluster <- function(nslaves) {
    if (!require(Rmpi))
        stop("Requires RMPI package.")
    # first, spawn nslaves processes on as many nodes as are
    # available
    comm <- 1
    mpi.spawn.Rslaves(nslaves=nslaves)
    cat("Slaves spawned successfully.\n")

    # have each slave do a repeat loop where they are constantly
    # querying for messages from the master
    mpi.remote.exec(source("commonFuncs.R"))
    mpi.remote.exec(source("source.R"))
    mpi.remote.exec(recvQuery(master), ret=FALSE)    

    # aggregate results from the slaves and save into a list
    cl <- vector("list", nslaves)
    for (i in seq(along=cl)) cl[[i]] <- newSlave(i,comm)
    cl
}

stopCluster <- function() {
    #TODO: check if slaves have been spawned first
    mpi.close.Rslaves() #this seems to put R into a deadlock (??)
}

###################################################################
# Higher level functions
#

# machinery for task pull
taskPull <- function(ens_size, cl, fun, n, args) {
    p <- length(cl)
    if (n>0 && p>0) {
    
        #lambda function to submit a job to a node
        submit <- function(node,job)
            sendCall(cl[[node]], fun, args(ens_size,job), tag=job)

        #submit as many jobs as there are nodes (less if n<p)
        cat("Submitting jobs..........\n")
        for (i in 1:min(n,p)) submit(i,i)

        #receive jobs as they complete, saving results
        cat("Jobs sent. Receiving jobs.......\n")
        val <- vector("list", n)
        for (i in 1:n) {
            # wait to receive a result
            d <- recvIResult(cl)
#            browser()
            # j<=n => there are more tasks. submit now that
            # one node is free
            if (i+min(n,p)<=n) submit(d$node, i+min(n,p))

            # write the recently obtained result where it
            # belongs
            val[d$tag] <- list(d$value)
        }
        # return
        val
    }
}

# interface for regular task pulling
# x is a list where each item constitutes the dataset that each node will
# operate on
taskPullNormal <- function(cl, x, fun, ...) {
    cat("Beginning task pulling...\n")
    args <- function(ens_size,i) c(list(x[[i]]), list(...))
    taskPull(ens_size=0, cl, fun, length(x), args)
}

# interface for ensemble task pulling
# x is a matrix, subsets of which are sampled by each slave node and
# operated on
# can be used to perform ensemble learning by ML algorithms, or cross-
# validation which requires testing on various subsets
taskPullEnsemble <- function(ens_runs, ens_size, cl, x, fun, ...) {
    cat("Beginning task pulling...\n")

    # prepare function arguments
    args <- function(ens_size, i) {
        # random sample of size ens_size from our dataset
        nobs <- dim(x[[1]])[1]
        if (ens_size == nobs) {
            #we're operating on full dataset, likely don't want randomization
            rand_idx <- 1:nobs
        } 
        else {
            # randomize
            rand_idx <- sample(1:nobs,ens_size)
        }
        c(list(x[[1]][rand_idx,]), list(...)) 
    }

    # pull tasks to slaves
    taskPull(ens_size, cl, fun, ens_runs, args)
}

